#include "RKHumanAttr.h"



bool RKHumanAttr::Init(std::string modelPath)
{
    auto option = fastdeploy::RuntimeOption();
    option.UseRKNPU2();
    auto format = fastdeploy::ModelFormat::RKNN;

    std::string ageGenderModelPath = modelPath + "/age_gender";
    auto ag_model_file = ageGenderModelPath + "/model_fix_rk3588.rknn";
    auto ag_params_file = "";
    auto ag_config_file = ageGenderModelPath + "/infer_cfg.yml";



    mAgeGenderModel = std::make_shared<fastdeploy::vision::classification::PaddleClasModel>(ag_model_file, ag_params_file, ag_config_file, option, format);
    if (!mAgeGenderModel->Initialized()) {
      std::cerr << "Failed to initialize AgeGender Model." << std::endl;
      return false;
    }
    mAgeGenderModel->GetPreprocessor().DisablePermute();
    

    std::string appearanceModelPath = modelPath + "/appearance";
    auto ap_model_file = appearanceModelPath + "/model_fix_rk3588.rknn";
    auto ap_params_file = "";
    auto ap_config_file = appearanceModelPath + "/infer_cfg.yml";



    mAppearanceModel = std::make_shared<fastdeploy::vision::detection::PicoDet>(ap_model_file, ap_params_file, ap_config_file, option, format);  

    mAppearanceModel->GetPreprocessor().DisablePermute();
    mAppearanceModel->GetPreprocessor().DisableNormalize();
    mAppearanceModel->GetPostprocessor().ApplyNMS();



}

std::shared_ptr<AttrResult> RKHumanAttr::Infer(cv::Mat& image)
{
    fastdeploy::vision::DetectionResult ap_res;
    if (!mAppearanceModel->Predict(&image, &ap_res)) {
        std::cerr << "Failed to predict Appearance." << std::endl;
        return nullptr;
    }

    
    
    fastdeploy::vision::ClassifyResult ag_res;
    if (!mAgeGenderModel->Predict(&image, &ag_res,50)) {
      std::cerr << "Failed to predict AgeGender." << std::endl;
      return nullptr;
    }


    std::shared_ptr<AttrResult> sp_Ret = std::make_shared<AttrResult>();

    std::map<int,float> attrs;
    for(int i = 0 ; i < ag_res.label_ids.size();i++)
    {
        attrs[ag_res.label_ids[i]] = ag_res.scores[i];
    }

    if(attrs[22] > 0.5)
    {
      sp_Ret->Age = 0;
    }
    else
    {
      sp_Ret->Age = 1;
    }
    std::vector<float> ageScores = { attrs[19] , attrs[20] , attrs[21] };

    sp_Ret->Gender = std::max_element(ageScores.begin(),ageScores.end()) - ageScores.begin();

    for(int i = 0 ; i < ap_res.label_ids.size();i++)
    {
        sp_Ret->Appearance.push_back(ap_res.label_ids[i]);
    }

    printf("\n Gender : %d",sp_Ret->Gender);
    printf("\n Age : %d",sp_Ret->Age);
    for(int i = 0 ; i < sp_Ret->Appearance.size();i++)
    {
        printf("\n Appearance : %d",sp_Ret->Appearance[i]);
    }
    printf("\n");
    return sp_Ret;

}